from typing import TYPE_CHECKING

from zkl_training import ProcessingTaskPlugin

if TYPE_CHECKING:
    from .training import GPTTraining


class CheckpointPlugin(ProcessingTaskPlugin):
    def __init__(self):
        super().__init__()
        self.last_checkpoint_tokens_n: int | None = None

    @property
    def task(self) -> 'GPTTraining':
        task = super().task
        from .training import GPTTraining
        assert isinstance(task, GPTTraining)
        return task

    @property
    def progress_tokens_n(self) -> int:
        return self.task.progress_tokens_n

    # task

    def on_before_run(self):
        self.last_checkpoint_tokens_n = self.progress_tokens_n

    def on_after_step(self):
        if self.check_need_checkpoint():
            self.checkpoint()
            self.last_checkpoint_tokens_n = self.progress_tokens_n

    # summarizing

    def check_need_checkpoint(self):
        checkpoint_tokens_n = self.task.training_hparams.summary_tokens_n
        if checkpoint_tokens_n is None:
            return False

        if checkpoint_tokens_n <= 0:
            return True

        warmup_tokens_n = self.task.training_hparams.warmup_tokens_n
        if warmup_tokens_n is not None:
            if self.last_checkpoint_tokens_n < warmup_tokens_n:
                checkpoint_tokens_n //= 10

        next_checkpoint_tokens_n = ((self.last_checkpoint_tokens_n // checkpoint_tokens_n) + 1) * checkpoint_tokens_n
        return self.progress_tokens_n >= next_checkpoint_tokens_n

    def checkpoint(self):
        if self.task.training_dir_path is not None:
            self.task.save_checkpoint()
